import time

import seaborn as sns
import matplotlib.pyplot as plt

from tensorboard.backend.event_processing import event_accumulator

import os
from queue import Queue

import pandas as pd


def get_tensorboard_info(log_path, key='score'):
    ea = event_accumulator.EventAccumulator(log_path)
    ea.Reload()

    val_psnr = ea.scalars.Items(key)
    x = []
    y = []
    for i in val_psnr:
        x.append(i.step)
        y.append(i.value)

    return x, y


def get_log_paths(dir_path):
    dir_queue = Queue()
    log_queue = Queue()
    if os.path.isdir(dir_path):
        dir_queue.put(dir_path)
    else:
        print(f'传入的参数{dir_path}不是一个目录，函数结束')

    while True:
        current_path = dir_queue.get()
        for path in os.listdir(current_path):
            if os.path.isdir(os.path.join(current_path, path)):
                dir_queue.put(os.path.join(current_path, path))
            else:
                if path.startswith('events'):
                    log_queue.put(os.path.join(current_path, path))
                else:
                    pass

        if dir_queue.empty():
            break

    return log_queue


def plot_info(x_info, y_info, plt_ax, smooth_factor=0.2, label=None):
    # 加载日志数据

    if label is not None:
        plt_ax.plot(pd.Series(y_info, index=x_info).ewm(alpha=smooth_factor).mean(), label=label)
    else:
        plt_ax.plot(pd.Series(y_info, index=x_info).ewm(alpha=smooth_factor).mean())


def plot_multi_lines(paths, labels, x_label, y_label, title):
    fig, ax = plt.subplots(1, 1)
    for path, label in zip(paths, labels):
        x, y = get_tensorboard_info(path)
        plot_info(x, y, ax, label=label)
    ax.legend()
    ax.set_xlabel(x_label)
    ax.set_ylabel(y_label)
    ax.set_title(title)
    fig.show()


if __name__ == '__main__':
    sns.set_style('darkgrid')

    # 获取目标文件夹下的tensorboard日志文件
    log_path = 'ablation_exps/Surviving/easy/withoutMulti'

    a = get_log_paths(log_path)

    path: str = ''
    path_list: list = []
    all_x = []
    all_y = []
    all_type = []
    csv_pdf_name = 'easy-multi'

    for model in os.listdir(log_path):
        for path in list(a.queue):
            if model in path:
                path_list.append(path)

        for path in path_list:
            x, y = get_tensorboard_info(path)
            all_x += x
            all_y += y
            all_type += [model] * len(x)
        path_list.clear()

    my_dict = {'score': all_y, 'time_stamp': all_x, 'model_type': all_type}
    df = pd.DataFrame(my_dict)
    df.to_csv(f'./dataset/{csv_pdf_name}.csv')
    # df = pd.read_csv('./dataset/hard.csv')

    image = sns.relplot(data=df, x='time_stamp', y='score', hue='model_type', kind='line')

    # image.fig.subplots_adjust(top=0.9)
    # image.fig.suptitle('HELLO')

    image.savefig(f'./dataset/{csv_pdf_name}.pdf')
